

package org.tgit.gateway.filter;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.PrincipalNameKeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.tgit.gateway.common.ApiConstants;
import org.tgit.gateway.common.TgitResponse;

import reactor.core.publisher.Flux;


import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.nio.ByteBuffer;
import java.util.Map;


/**
 * 
* @ClassName: RequestRateLimiterGatewayFilterFactory   
* @Description:  
* @author fanzaijun 
* @date 2018年11月6日 下午3:00:10   
*
 */
@Component
public class NetLimitGatewayFilterFactory extends AbstractGatewayFilterFactory<NetLimitGatewayFilterFactory.Config> {

	//限流方式自定义
	@Value("${limit.case}") 
	private String limitCase;
	
	public static final String KEY_RESOLVER_KEY = "keyResolver";
	
	@SuppressWarnings("rawtypes")
	private final RateLimiter defaultRateLimiter;
	private final KeyResolver defaultKeyResolver;
	
	//这里是PrincipalNameKeyResolver.源码有说明。
	@SuppressWarnings("rawtypes")
	public NetLimitGatewayFilterFactory(RateLimiter defaultRateLimiter, PrincipalNameKeyResolver defaultKeyResolver) {
		super(Config.class);
		this.defaultRateLimiter = defaultRateLimiter;
		this.defaultKeyResolver = defaultKeyResolver;
	}

	public KeyResolver getDefaultKeyResolver() {
		return defaultKeyResolver;
	}

	@SuppressWarnings("rawtypes")
	public RateLimiter getDefaultRateLimiter() {
		return defaultRateLimiter;
	}

	@SuppressWarnings({ "unchecked", "unused" })
	@Override
	public GatewayFilter apply(Config config) {
		KeyResolver resolver = (config.keyResolver == null) ? defaultKeyResolver : config.keyResolver;
		RateLimiter<Object> limiter = (config.rateLimiter == null) ? defaultRateLimiter : config.rateLimiter;

		return (exchange, chain) -> {
			Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
			TgitResponse tgitResponse=new TgitResponse();
			//113,并发限流错误
	    	tgitResponse.setMessage(ApiConstants.ERROR_MAP.get(ApiConstants.REQ_ERROR_CODE_113));
			ClientResponse clientResponse = exchange.getAttribute(ServerWebExchangeUtils.CLIENT_RESPONSE_ATTR);
			return  resolver.resolve(exchange).flatMap(key ->
					// TODO: if key is empty?
					limiter.isAllowed(route.getId(), key).flatMap(response -> {
		
						for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
							exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
						}
		
						if (response.isAllowed()) {
							return chain.filter(exchange);
						}
			
						//exchange.getResponse().setStatusCode(config.getStatusCode());
						exchange.getResponse().writeWith(Flux.just(this.getBodyBuffer(exchange.getResponse(),tgitResponse)));
						Route newRoute=Route.async()
								.asyncPredicate(route.getPredicate())
								.filters(route.getFilters())
								.id(route.getId())
								.order(route.getOrder())
								.uri("forward:///rest/gateway/manager/error/113")
								.build();
						exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR, newRoute);
					    return chain.filter(exchange);
		         }));
		};
	}

	@SuppressWarnings("rawtypes")
	public static class Config {
		private KeyResolver keyResolver;
		private RateLimiter rateLimiter;
		private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;

		public KeyResolver getKeyResolver() {
			return keyResolver;
		}

		public Config setKeyResolver(KeyResolver keyResolver) {
			this.keyResolver = keyResolver;
			return this;
		}
		public RateLimiter getRateLimiter() {
			return rateLimiter;
		}

		public Config setRateLimiter(RateLimiter rateLimiter) {
			this.rateLimiter = rateLimiter;
			return this;
		}

		public HttpStatus getStatusCode() {
			return statusCode;
		}

		public Config setStatusCode(HttpStatus statusCode) {
			this.statusCode = statusCode;
			return this;
		}
	}

	private DataBuffer getBodyBuffer(ServerHttpResponse response, TgitResponse result) {
        try {
			return response.bufferFactory().wrap(ByteBuffer.wrap(this.getBytes(result)));
		} catch (IOException e) {
			e.printStackTrace();
		}
		return null;
    }
	
	private  byte[] getBytes(Serializable obj) throws IOException {
        ByteArrayOutputStream bout = new ByteArrayOutputStream();
        ObjectOutputStream out = new ObjectOutputStream(bout);
        out.writeObject(obj);
        out.flush();
        byte[] bytes = bout.toByteArray();
        bout.close();
        out.close();
        return bytes;
    }
}
